﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using CPAIS.Updater;

namespace CPAIS.Algorithm
{

/**
 * Hierarchical Dirichlet Processes  
 * Chinese Restaurant Franchise Sampler
 * 
 */
    public class HDPGibbs
    {

        public double beta = 0.5; // default only
        public double gamma = 1.5;
        public double alpha = 1.0;

        private Random random = new Random();
        private double[] p;
        private double[] f;

        public DOCState[] docStates;
        public int[] numberOfTablesByTopic;
        public int[] wordCountByTopic;
        public int[][] wordCountByTopicAndTerm;


        public int sizeOfVocabulary;
        public int totalNumberOfWords;
        public int numberOfTopics = 1;
        public int totalNumberOfTables;


        /**
         * Initially assign the words to tables and topics
         * 
         * @param corpus {@link CLDACorpus} on which to fit the model
         */
        public void addInstances(int[][] documentsInput, int V)
        {
            sizeOfVocabulary = V;
            totalNumberOfWords = 0;
            docStates = new DOCState[documentsInput.Length];
            for (int d = 0; d < documentsInput.Length; d++)
            {
                docStates[d] = new DOCState(documentsInput[d], d);
                totalNumberOfWords += documentsInput[d].Length;
            }
            int k, i, j;
            DOCState docState;
            p = new double[20];
            f = new double[20];
            numberOfTablesByTopic = new int[numberOfTopics + 1];
            wordCountByTopic = new int[numberOfTopics + 1];
            wordCountByTopicAndTerm = new int[numberOfTopics + 1][];
            for (k = 0; k <= numberOfTopics; k++) 	// var initialization done
                wordCountByTopicAndTerm[k] = new int[sizeOfVocabulary];
            for (k = 0; k < numberOfTopics; k++)
            {
                docState = docStates[k];
                for (i = 0; i < docState.documentLength; i++)
                    addWord(docState.docID, i, 0, k);
            } // all topics have now one document
            for (j = numberOfTopics; j < docStates.Length; j++)
            {
                docState = docStates[j];
                k = random.Next(numberOfTopics);
                for (i = 0; i < docState.documentLength; i++)
                    addWord(docState.docID, i, 0, k);
            } // the words in the remaining documents are now assigned too
        }

        public DOCState getInstance(int nDocIndex)
        {
            return this.docStates[nDocIndex];
        }
        /**
         * Step one step ahead
         * 
         */
        protected void nextGibbsSweep()
        {
            int table;
            for (int d = 0; d < docStates.Length; d++)
            {
                for (int i = 0; i < docStates[d].documentLength; i++)
                {
                    removeWord(d, i); // remove the word i from the state
                    table = sampleTable(d, i);
                    if (table == docStates[d].numberOfTables) // new Table
                        addWord(d, i, table, sampleTopic()); // sampling its Topic
                    else
                        addWord(d, i, table, docStates[d].tableToTopic[table]); // existing Table
                }
            }
            defragment();
        }


        /**
         * Decide at which topic the table should be assigned to
         * 
         * @return the index of the topic
         */
        private int sampleTopic()
        {
            double u, pSum = 0.0;
            int k;
            p = ensureCapacity(p, numberOfTopics);
            for (k = 0; k < numberOfTopics; k++)
            {
                pSum += numberOfTablesByTopic[k] * f[k];
                p[k] = pSum;
            }
            pSum += gamma / sizeOfVocabulary;
            p[numberOfTopics] = pSum;
            u = random.NextDouble() * pSum;
            for (k = 0; k <= numberOfTopics; k++)
                if (u < p[k])
                    break;
            return k;
        }


        /**	 
         * Decide at which table the word should be assigned to
         * 
         * @param docID the index of the document of the current word
         * @param i the index of the current word
         * @return the index of the table
         */
        int sampleTable(int docID, int i)
        {
            int k, j;
            double pSum = 0.0, vb = sizeOfVocabulary * beta, fNew, u;
            DOCState docState = docStates[docID];
            f = ensureCapacity(f, numberOfTopics);
            p = ensureCapacity(p, docState.numberOfTables);
            fNew = gamma / sizeOfVocabulary;

            for (k = 0; k < numberOfTopics; k++)
            {
                f[k] = (wordCountByTopicAndTerm[k][docState.words[i].termIndex] + beta) /
                        (wordCountByTopic[k] + vb);
                fNew += numberOfTablesByTopic[k] * f[k];
            }

            for (j = 0; j < docState.numberOfTables; j++)
            {
                if (docState.wordCountByTable[j] > 0)
                    pSum += docState.wordCountByTable[j] * f[docState.tableToTopic[j]];
                p[j] = pSum;
            }
            pSum += alpha * fNew / (totalNumberOfTables + gamma); // Probability for t = tNew
            p[docState.numberOfTables] = pSum;
            u = random.NextDouble() * pSum;
            for (j = 0; j <= docState.numberOfTables; j++)
                if (u < p[j])
                    break;	// decided which table the word i is assigned to
            return j;
        }


        /**
         * Method to call for fitting the model.
         * 
         * @param doShuffle
         * @param shuffleLag
         * @param maxIter number of iterations to run
         * @param saveLag save interval 
         * @param wordAssignmentsWriter {@link WordAssignmentsWriter}
         * @param topicsWriter {@link TopicsWriter}
         * @throws IOException 
         */
        public void run(int shuffleLag, int maxIter, ProgressUpdater updater)
        {
            if (updater != null)
                updater.UpdateMessage("Initializing...");
            if (updater != null)
                updater.UpdateMessage("Gibbs sampling...");

            for (int iter = 0; iter < maxIter; iter++)
            {
                if ((shuffleLag > 0) && (iter > 0) && (iter % shuffleLag == 0))
                    doShuffle();
                nextGibbsSweep();
                Console.WriteLine("iter = " + iter + " #topics = " + numberOfTopics + ", #tables = "
                        + totalNumberOfTables);

                if (updater != null)
                    updater.UpdateProgress((double)(iter + 1) / maxIter);
            }

            if (updater != null)
                updater.UpdateMessage("Calcaulating result...");
        }


        /**
         * Removes a word from the bookkeeping
         * 
         * @param docID the id of the document the word belongs to 
         * @param i the index of the word
         */
        protected void removeWord(int docID, int i)
        {
            DOCState docState = docStates[docID];
            int table = docState.words[i].tableAssignment;
            int k = docState.tableToTopic[table];
            docState.wordCountByTable[table]--;
            wordCountByTopic[k]--;
            wordCountByTopicAndTerm[k][docState.words[i].termIndex]--;
            if (docState.wordCountByTable[table] == 0)
            {
                // table is removed
                totalNumberOfTables--;
                numberOfTablesByTopic[k]--;
                docState.tableToTopic[table]--;
            }
        }



        /**
         * Add a word to the bookkeeping
         * 
         * @param docID	docID the id of the document the word belongs to 
         * @param i the index of the word
         * @param table the table to which the word is assigned to
         * @param k the topic to which the word is assigned to
         */
        protected void addWord(int docID, int i, int table, int k)
        {
            DOCState docState = docStates[docID];
            docState.words[i].tableAssignment = table;
            docState.wordCountByTable[table]++;
            wordCountByTopic[k]++;
            wordCountByTopicAndTerm[k][docState.words[i].termIndex]++;
            if (docState.wordCountByTable[table] == 1)
            {
                // a new table is created
                docState.numberOfTables++;
                docState.tableToTopic[table] = k;
                totalNumberOfTables++;
                numberOfTablesByTopic[k]++;
                docState.tableToTopic = ensureCapacity(docState.tableToTopic, docState.numberOfTables);
                docState.wordCountByTable = ensureCapacity(docState.wordCountByTable, docState.numberOfTables);
                if (k == numberOfTopics)
                {
                    // a new topic is created
                    numberOfTopics++;
                    numberOfTablesByTopic = ensureCapacity(numberOfTablesByTopic, numberOfTopics);
                    wordCountByTopic = ensureCapacity(wordCountByTopic, numberOfTopics);
                    wordCountByTopicAndTerm = add(wordCountByTopicAndTerm, new int[sizeOfVocabulary], numberOfTopics);
                }
            }
        }


        /**
         * Removes topics from the bookkeeping that have no words assigned to
         */
        protected void defragment()
        {
            int[] kOldToKNew = new int[numberOfTopics];
            int k, newNumberOfTopics = 0;
            for (k = 0; k < numberOfTopics; k++)
            {
                if (wordCountByTopic[k] > 0)
                {
                    kOldToKNew[k] = newNumberOfTopics;
                    swap(wordCountByTopic, newNumberOfTopics, k);
                    swap(numberOfTablesByTopic, newNumberOfTopics, k);
                    swap(wordCountByTopicAndTerm, newNumberOfTopics, k);
                    newNumberOfTopics++;
                }
            }
            numberOfTopics = newNumberOfTopics;
            for (int j = 0; j < docStates.Length; j++)
                docStates[j].defragment(kOldToKNew);
        }


        /**
         * Permute the ordering of documents and words in the bookkeeping
         */
        protected void doShuffle()
        {
            List<DOCState> h = docStates.ToList();
            List<DOCState> newList = new List<DOCState>();
            Random _rnd = new Random();
            //while there are still items in our list
            while (h.Count > 0)
            {
                //get a random number within the list
                int index = _rnd.Next(h.Count);

                //add the item at that position to the new list
                newList.Add(h[index]);

                //and remove it from the old list
                h.RemoveAt(index);
            }


            for (int j = 0; j < newList.Count; j++)
            {
                List<WordState> h2 = newList[j].words.ToList();

                List<WordState> w2 = new List<WordState>();
                while (h2.Count > 0)
                {
                    int index2 = _rnd.Next(h2.Count);
                    w2.Add(h2[index2]);
                    h2.RemoveAt(index2);
                }

                newList[j].words = w2.ToArray();
            }
            docStates = newList.ToArray();
        }


        public static void swap(int[] arr, int arg1, int arg2)
        {
            int t = arr[arg1];
            arr[arg1] = arr[arg2];
            arr[arg2] = t;
        }

        public static void swap(int[][] arr, int arg1, int arg2)
        {
            int[] t = arr[arg1];
            arr[arg1] = arr[arg2];
            arr[arg2] = t;
        }

        public static double[] ensureCapacity(double[] arr, int min)
        {
            int length = arr.Length;
            if (min < length)
                return arr;
            double[] arr2 = new double[min * 2];
            for (int i = 0; i < length; i++)
                arr2[i] = arr[i];
            return arr2;
        }

        public static int[] ensureCapacity(int[] arr, int min)
        {
            int length = arr.Length;
            if (min < length)
                return arr;
            int[] arr2 = new int[min * 2];
            for (int i = 0; i < length; i++)
                arr2[i] = arr[i];
            return arr2;
        }

        public static int[][] add(int[][] arr, int[] newElement, int index)
        {
            int length = arr.Length;
            if (length <= index)
            {
                int[][] arr2 = new int[index * 2][];
                for (int i = 0; i < length; i++)
                    arr2[i] = arr[i];
                arr = arr2;
            }
            arr[index] = newElement;
            return arr;
        }


        public class DOCState
        {
            public int docID, documentLength, numberOfTables;
            public int[] tableToTopic;
            public int[] wordCountByTable;
            public WordState[] words;

            public DOCState(int[] instance, int docID)
            {
                this.docID = docID;
                numberOfTables = 0;
                documentLength = instance.Length;
                words = new WordState[documentLength];
                wordCountByTable = new int[2];
                tableToTopic = new int[2];
                for (int position = 0; position < documentLength; position++)
                    words[position] = new WordState(instance[position], -1);
            }


            public void defragment(int[] kOldToKNew)
            {
                int[] tOldToTNew = new int[numberOfTables];
                int t, newNumberOfTables = 0;
                for (t = 0; t < numberOfTables; t++)
                {
                    if (wordCountByTable[t] > 0)
                    {
                        tOldToTNew[t] = newNumberOfTables;
                        tableToTopic[newNumberOfTables] = kOldToKNew[tableToTopic[t]];
                        swap(wordCountByTable, newNumberOfTables, t);
                        newNumberOfTables++;
                    }
                    else
                        tableToTopic[t] = -1;
                }
                numberOfTables = newNumberOfTables;
                for (int i = 0; i < documentLength; i++)
                    words[i].tableAssignment = tOldToTNew[words[i].tableAssignment];
            }

        }


        public class WordState
        {
            public int termIndex { set; get; }
            public int tableAssignment { set; get; }

            public WordState(int wordIndex, int tableAssignment)
            {
                this.termIndex = wordIndex;
                this.tableAssignment = tableAssignment;
            }

        }

        /**
        public static void main(String[] args) throws IOException 
         * {

            HDPGibbsSampler hdp = new HDPGibbsSampler();
            CLDACorpus corpus = new CLDACorpus(new FileInputStream(args[0]));
            hdp.addInstances(corpus.getDocuments(), corpus.getVocabularySize());

            System.out.println("sizeOfVocabulary = "+hdp.sizeOfVocabulary);
            System.out.println("totalNumberOfWords = "+hdp.totalNumberOfWords);
            System.out.println("NumberOfDocs = "+hdp.docStates.length);

            hdp.run(0, 2000, System.out);
		

            PrintStream file = new PrintStream(args[1]);
            for (int k = 0; k < hdp.numberOfTopics; k++) {
                for (int w = 0; w < hdp.sizeOfVocabulary; w++)
                    file.format("%05d ",hdp.wordCountByTopicAndTerm[k][w]);
                file.println();
            }
            file.close();

		
            file = new PrintStream(args[2]);
            file.println("d w z t");
            int t, docID;
            for (int d = 0; d < hdp.docStates.length; d++) {
                DOCState docState = hdp.docStates[d];
                docID = docState.docID;
                for (int i = 0; i < docState.documentLength; i++) {
                    t = docState.words[i].tableAssignment;
                    file.println(docID + " " + docState.words[i].termIndex + " " + docState.tableToTopic[t] + " " + t); 
                }
            }
            file.close();
		
        }
            **/
    }
}
